package main

import "fmt"

type ListNode struct {
	Val  int
	Next *ListNode
}

func mergeKLists(lists []*ListNode) *ListNode {
	//fmt.Println(lists)
	if len(lists) == 1 {
		if lists[0] != nil {
			return lists[0]
		}
		return nil
	}
	//虚拟头结点
	dummy := &ListNode{
		Val:  -1,
		Next: nil,
	}
	p := dummy
	//构建优先级队列，最小堆
	minpq := InitMinPQ(5)
	for _, item := range lists {
		if item == nil {
			continue
		}
		node := item
		for {
			minpq.Insert(node.Val)
			if node.Next == nil {
				break
			}
			node = node.Next
		}
	}
	//fmt.Println(minpq)
	len := minpq.size
	for i := 0; i < len; i++ {
		//fmt.Println(minpq.DelMin())
		p.Next = &ListNode{
			Val:  minpq.DelMin(),
			Next: nil,
		}
		p = p.Next
	}
	//fmt.Println(*dummy)
	return dummy.Next
}

type MinPQ struct {
	//save element
	pq []int
	//the number of pq element(no包括前置0)
	size int
}

func InitMinPQ(cap int) *MinPQ {
	return &MinPQ{
		pq:   make([]int, 1, cap+1),
		size: 0,
	}
}

func (p *MinPQ) parent(x int) int {
	return x / 2
}

func (p *MinPQ) left(x int) int {
	return x * 2
}

func (p *MinPQ) right(x int) int {
	return x*2 + 1
}

// MinVal MaxVal return MAX value of MaxPQ
func (p *MinPQ) MinVal() int {
	//fmt.Println(len(p.pq), cap(p.pq))
	if p.size < 1 {
		panic("no elements in MaxPQ")
	}
	return p.pq[1]
}

//Insert to insert an element
func (p *MinPQ) Insert(e int) {
	p.pq = append(p.pq, e)
	p.size += 1

	p.swim(p.size)
	//fmt.Println(p)
}

// DelMin DelMax 删除并返回最大值
func (p *MinPQ) DelMin() (min int) {
	min = p.pq[1]
	p.swap(1, p.size)
	p.pq = append(p.pq[:len(p.pq)-1], p.pq[len(p.pq):]...)
	p.size -= 1

	p.sink(1)
	return
}

//swim 上浮第x个元素
func (p *MinPQ) swim(x int) {
	node := x
	for {
		//如果是根节点 或者 父节点值小于节点值
		if node == 1 || p.less(p.parent(node), node) {
			//停止上浮
			break
		}
		//开始交换
		p.swap(node, p.parent(node))
		node = p.parent(node)
	}
}

//sink 下沉第x个元素
func (p *MinPQ) sink(x int) {
	node := x
	smaller := p.left(node)
	for {
		//沉到底了就不再沉了
		if smaller > p.size {
			break
		}
		//如果右节点存在并且不大于左结点 那么就和右节点交换
		if p.right(node) <= p.size && p.less(p.right(node), smaller) {
			smaller = p.right(node)
		}
		//如果子节点中最大的比自己还小 就不用下沉了
		if p.less(node, smaller) {
			break
		}
		//否则 继续下沉
		p.swap(smaller, node)
		node = smaller
		smaller = p.left(node)
	}
}

//swap 交换两个元素
func (p *MinPQ) swap(i, j int) {
	temp := p.pq[i]
	p.pq[i] = p.pq[j]
	p.pq[j] = temp
}

//less 比较两个元素大小
func (p *MinPQ) less(i, j int) bool {
	return p.pq[i] < p.pq[j]
}

func main() {
	lists := []*ListNode{
		&ListNode{
			Val: 1,
			Next: &ListNode{
				Val: 4,
				Next: &ListNode{
					Val:  5,
					Next: nil,
				},
			},
		},
		&ListNode{
			Val: 1,
			Next: &ListNode{
				Val: 3,
				Next: &ListNode{
					Val:  4,
					Next: nil,
				},
			},
		},
		&ListNode{
			Val: 2,
			Next: &ListNode{
				Val:  6,
				Next: nil,
			},
		},
	}
	res := mergeKLists(lists)
	fmt.Println(res)

	var lists2 []*ListNode
	lists2 = make([]*ListNode, 1)
	res2 := mergeKLists(lists2)
	fmt.Println(res2)

	var lists3 []*ListNode
	lists3 = make([]*ListNode, 1)
	lists3[0] = &ListNode{Val: 1, Next: nil}
	res3 := mergeKLists(lists3)
	fmt.Println(res3)

	var lists4 []*ListNode
	lists4 = make([]*ListNode, 2)
	res4 := mergeKLists(lists4)
	fmt.Println(res4)
}
